% clear; clc;
% 
% I = im2double(imread('cameraman.tif'));
% params.a = 1;
% params.b = 2*pi;
% params.c = 11.5;
% params.x0 = 0.4;
% params.y0 = 0.8;
% params.z0 = 0.6;
% params.m = 100;
% params.n = 500;
% params.CR = 0.8;
% 
% [C, keys] = encrypt_image(I, params);
% 
% figure;
% subplot(1,2,1);
% imshow(I);
% title('Original Image');
% 
% subplot(1,2,2);
% imshow(C);
% title('Encrypted Image');
% 
% function [C, keys] = encrypt_image(I, params)
%     % 输入：
%     % I      - 原始灰度图像，double类型，范围[0,1]
%     % params - 结构体参数，包含3D-SIMM参数，迭代次数，压缩率等
%     %
%     % 输出：
%     % C      - 加密图像 uint8矩阵
%     % keys   - 结构体，保存密钥参数，供解密使用
% 
%     % Step 1: DWT稀疏化，置零细节系数，恢复同尺寸稀疏矩阵
%     [cA, cH, cV, cD] = dwt2(I, 'haar');
%     cH(:) = 0; cV(:) = 0; cD(:) = 0;
%     I1 = idwt2(cA, cH, cV, cD, 'haar');
%     I1_vec = I1(:);
% 
%     % 补齐长度为2的幂
%     len_orig = length(I1_vec);
%     len = 2^nextpow2(len_orig);
%     if len > len_orig
%         I1_vec(len_orig+1:len) = 0;
%     end
% 
%     % 计算采样点数M和边长side_len（完全平方数）
%     M_approx = floor(len * params.CR);
%     side_len = floor(sqrt(M_approx));
%     M = side_len^2;
%     if M == 0
%         error('采样点数 M 计算为0，压缩率过低或图像尺寸太小');
%     end
% 
%     % Step 2: 生成混沌序列X，长度m+M，取后M
%     len_X = params.m + M;
%     [X, ~, ~] = gen_3DSIMM_seq(params.a, params.b, params.c, ...
%                                params.x0, params.y0, params.z0, len_X);
%     X = X(params.m+1:end);
% 
%     % Step 3: 根据X排序得到采样行索引
%     [~, idx] = sort(X, 'descend');
%     selected_rows = idx(1:M);
% 
%     % Step 4: 快速Walsh-Hadamard采样
%     I2_vec = fwht_sampling(I1_vec, selected_rows);
% 
%     % Step 5: 重塑为二维压缩图像
%     I2 = reshape(I2_vec(1:M), side_len, side_len);
% 
%     % Step 6: 归一化量化到[0,255]
%     I_min = min(I2(:));
%     I_max = max(I2(:));
%     if I_max == I_min
%         error('归一化失败，最大值等于最小值');
%     end
%     I2_norm = (I2 - I_min) / (I_max - I_min);
%     I2q = uint8(floor(I2_norm * 255));
% 
%     % Step 7: 生成置乱序列，长度m+M*M，取后部分
%     len_scramble = params.m + M;
%     [X_scr, Y_scr, ~] = gen_3DSIMM_seq(params.a, params.b, params.c, ...
%                                        params.x0, params.y0, params.z0, len_scramble);
%     X_scr = X_scr(params.m+1:end);
%     Y_scr = Y_scr(params.m+1:end);
% 
%     a_scr = floor(abs(X_scr(1) * 1e6));
%     b_scr = floor(abs(Y_scr(1) * 1e6));
% 
%     % Step 8: Arnold二维置乱（示范用10次迭代）
%     I3 = arnold_2d_scramble(I2q, 10);
% 
%     % Step 9: 矩阵转置
%     I3 = I3';
% 
%     % Step 10: 生成扩散密钥序列
%     len_diff = params.n + M;
%     [~, y_diff, z_diff] = gen_3DSIMM_seq(params.a, params.b, params.c, ...
%                                          params.x0, params.y0, params.z0, len_diff);
%     y_diff = y_diff(params.n+1:end);
%     z_diff = z_diff(params.n+1:end);
% 
%     S1 = mod(floor(y_diff * 2^10), 256);
%     S2 = mod(floor(z_diff * 2^10), 256);
% 
%     % Step 11: 双向扩散加密
%     C = diffusion_double_fixed(I3, S1, S2);
% 
%     % 返回密钥信息
%     keys.Phi_rows = selected_rows;
%     keys.I_min = I_min;
%     keys.I_max = I_max;
%     keys.side_len = side_len;
%     keys.a_scr = a_scr;
%     keys.b_scr = b_scr;
%     keys.S1 = S1;
%     keys.S2 = S2;
% end
% 
% %% 辅助函数
% 
% function y = fwht_sampling(x, rows)
%     N = length(x);
%     if bitand(N, N-1) ~= 0
%         error('输入向量长度必须是2的幂');
%     end
%     Hx = fwht(x) / sqrt(N);
%     y = Hx(rows);
% end
% 
% function [X, Y, Z] = gen_3DSIMM_seq(a, b, c, x0, y0, z0, len)
%     X = zeros(1,len); Y = zeros(1,len); Z = zeros(1,len);
%     X(1)=x0; Y(1)=y0; Z(1)=z0;
%     for i=1:len-1
%         [X(i+1), Y(i+1), Z(i+1)] = SIMM(X(i), Y(i), Z(i), a, b, c);
%     end
% end
% 
% function [x_next, y_next, z_next] = SIMM(x, y, z, a, b, c)
%     x_next = a*sin(b*z)*sin(c/x);
%     y_next = a*sin(b*x_next)*sin(c/y);
%     z_next = a*sin(b*y_next)*sin(c/z);
% end
% 
% function I3 = arnold_2d_scramble(I2q, iter)
%     [N, ~] = size(I2q);
%     I3 = I2q;
%     for k = 1:iter
%         temp = zeros(N,N,'uint8');
%         for x = 0:N-1
%             for y = 0:N-1
%                 new_x = mod(x + y, N);
%                 new_y = mod(x + 2*y, N);
%                 temp(new_x+1, new_y+1) = I3(x+1, y+1);
%             end
%         end
%         I3 = temp;
%     end
% end
% 
% function C = diffusion_double_fixed(I3, S1, S2)
%     [M, N] = size(I3);
%     len = M*N;
%     I_vec = reshape(I3, [], 1);
%     C_vec = zeros(len, 1, 'uint16');
%     
%     % 前向扩散
%     C_vec(1) = mod(double(I_vec(1)) + double(S1(1)), 256);
%     for i = 2:len
%         C_vec(i) = mod(C_vec(i-1) + double(I_vec(i)) + double(S1(i)), 256);
%     end
%     
%     % 逆向扩散
%     for i = len-1:-1:1
%         C_vec(i) = mod(C_vec(i) + C_vec(i+1) + double(S2(i)), 256);
%     end
%     
%     C = uint8(C_vec);
%     C = reshape(C, M, N);
% end

function I_dec = decrypt_image(C, keys, params)
    % Step 1: 逆扩散
    I3 = inverse_diffusion_double(C, keys.S1, keys.S2);
    
    % Step 2: 矩阵转置（解置乱前的转置）
    I3 = I3';
    
    % Step 3: 逆Arnold二维置乱
    I2q = inverse_arnold_2d_scramble(I3, 10);
    
    % Step 4: 归一化逆量化
    I2_norm = double(I2q) / 255;
    I2 = I2_norm * (keys.I_max - keys.I_min) + keys.I_min;
    
    % Step 5: 压缩感知重构（简单伪逆重构）
    % 这里用快速Walsh-Hadamard变换伪逆
    len = length(keys.Phi_rows) + length(I2q(:)) - length(keys.Phi_rows); % 估计len
    side_len = keys.side_len;
    M = side_len^2;
    
    % 生成Hadamard完整矩阵长度（2的幂）
    N = 2^nextpow2(len);
    
    % 快速Hadamard变换函数（FWHT）需自己实现或调用
    % 伪逆重构示例：
    % 1) 构造Phi对应行号
    rows = keys.Phi_rows;
    
    % 2) 构造Hadamard矩阵对应函数
    % 3) 用伪逆估计x
    
    % 由于复杂，示范用简单插值恢复（效果有限）
    I1_vec_est = simple_cs_reconstruction(I2(:), rows, N);
    
    % Step 6: 补齐长度调整并逆小波重构
    len_orig = keys.len_orig;
    coeffs = I1_vec_est(1:len_orig);
    
    % 使用原始参数idwt2逆变换，简单示范为近似重构
    % 由于原加密用的是置零细节，解密时还原完整
    % 这里为示范，假设cH,cV,cD为0
    level = 1;
    waveletName = 'haar';
    sizes = calc_wavedec_sizes(size(I), level); % 需要自己实现或保存
    
    I_dec = waverec2(coeffs, sizes, waveletName);
    I_dec = mat2gray(I_dec);
end

% 逆扩散函数示范
function I3 = inverse_diffusion_double(C, S1, S2)
    % 逆扩散过程的反向实现，根据论文算法反写
    [M, N] = size(C);
    len = M*N;
    C_vec = reshape(C, [], 1);
    
    % 先逆向扩散反演前向扩散
    tmp = zeros(len, 1, 'double');
    tmp(1) = mod(C_vec(1) - S1(1), 256);
    for i=2:len
        tmp(i) = mod(C_vec(i) - S1(i) - tmp(i-1), 256);
    end
    
    % 再前向扩散反演逆向扩散
    I3_vec = zeros(len, 1, 'double');
    I3_vec(len) = mod(tmp(len) - S2(len), 256);
    for i = len-1:-1:1
        I3_vec(i) = mod(tmp(i) - S2(i) - I3_vec(i+1), 256);
    end
    
    I3 = reshape(uint8(I3_vec), M, N);
end

% 逆Arnold二维置乱示范
function I2q = inverse_arnold_2d_scramble(I3, iter)
    % 逆Arnold变换迭代次数同加密时相同
    [N, ~] = size(I3);
    I2q = I3;
    for k = 1:iter
        temp = zeros(N,N,'uint8');
        for x = 0:N-1
            for y = 0:N-1
                orig_x = mod(2*x - y, N);
                orig_y = mod(-x + y, N);
                temp(orig_x+1, orig_y+1) = I2q(x+1, y+1);
            end
        end
        I2q = temp;
    end
end

% 简单CS重构示范（伪逆估计）
function x_est = simple_cs_reconstruction(y, rows, N)
    % 随机生成Hadamard矩阵部分重构，仅示范
    H = hadamard(N) / sqrt(N);
    Phi = H(rows, :);
    x_est = pinv(Phi) * y;
end

% 计算wavedec2尺寸结构（示范）
function sizes = calc_wavedec_sizes(imgSize, level)
    % 这里只示范同样大小，可根据需要实现准确计算
    N = imgSize(1);
    sizes = repmat([N N], level+1, 1);
end

main_encrypt_decrypt();